import torch.nn as nn

from utils import clone_module
from layer_norm import LayerNorm


class Decoder(nn.Module):

    def __init__(self, decoder_layer, num_decoder_layer):
        super(Decoder, self).__init__()
        self.decoder_layers = clone_module(decoder_layer, num_decoder_layer)
        self.norm = LayerNorm(decoder_layer.feature_dim)

    def forward(self, x, memory, source_mask, target_mask):
        for layer in self.decoder_layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)